package scu.maqiang.numeric;

import java.util.Arrays;

/**
 * 稠密矩阵直解求解器, 实现基本的列选主元LU分解与Choleskty分解�ֽⷨ
 * @author 马强
 *
 */
public class DirectDSolver extends Solver{
	
	/**
	 * 默认构造函数
	 */
	public DirectDSolver() {
	}
	
	/**
	 * 构造函数，参数为一个稠密方阵
	 * @param A 稠密矩阵，要求为方阵
	 */
	public DirectDSolver(double[][] A) {
		int n = A.length;
		for(int i = 0; i < n; i++) {
			if (A[i].length != n) {
				throw new IllegalArgumentException("The matrix is not square!");
			}
		}
		this.n = n;
		this.A = A;
	}
	
	/**
	 * 设置求解器关联的矩阵
	 * @param A 稠密矩阵，要求为方阵
	 */
	public void setMatrix(double[][] A) {
		int n = A.length;
		for(int i = 0; i < n; i++) {
			if (A[i].length != n) {
				throw new IllegalArgumentException("The matrix is not square!");
			}
		}
		this.n = n;
		this.A = A;	
	}
	
	/**
	 * 使用列选主元高斯消去除求解线性方程组, 直接搞用
	 * @param b 线性方程组右端项
	 * @param x 返回方程的解
	 */
	public void PGaussSolve(double[] b, double[] x) {
		PLU();
		PLUSolve(b, x);
	}

	/**
	 * 使用列选主元高斯消去除求解线性方程组, 直接搞用
	 * @param b 线性方程组右端项
	 * @return 返回方程的解
	 */
	public double[] PGaussSolve(double[] b) {
		PLU();
		return PLUSolve(b);
	}
	
	
	/**
	 * 对关联矩阵进行列选主元LU分解
	 */
	public void PLU() {
		P = new int[n];
		Arrays.setAll(P, i -> i);
		for(int i = 0; i < n - 1; i++) {
			Pair<Integer, Double> maxIdxResult = MVO.maxAbsColValDown(A, i);
			int idxJ = maxIdxResult.getFirst();
			double maxval = maxIdxResult.getSecond();
			if (Math.abs(maxval) < Constants.Er) {
				throw new ArithmeticException("Zero pivot element!");
			}
			if (idxJ != i) {
				double[] tempRow = A[i]; A[i] = A[idxJ]; A[idxJ] = tempRow;
				int tempIdx = P[i]; P[i] = P[idxJ]; P[idxJ] = tempIdx;
			}
			
			for(int j = i + 1; j < n; j++) {
				if (Math.abs(A[j][i]) > Constants.Er) {
					double ratio = A[j][i] / maxval;
					A[j][i] = ratio;
					for(int k = i + 1; k < n; k++) {
						A[j][k] -= A[i][k] * ratio;
					}
				}
			}
		}
	}
	
	/**
	 * 在对矩阵进行列选主元LU分解之上求解线性方程组, 使用该函数前需要使用PLU()函数对矩阵作分解
	 * @param b 线性方程组右端项
	 * @param x 方程组解
	 */
	public void PLUSolve(double[] b, double[] x) {
		Arrays.setAll(x, i -> b[P[i]]);
		double temp = 0.0;
		for(int i = 0; i < n; i++) {
			temp = 0.0;
			for(int j = 0; j < i; j++) {
				temp += A[i][j] * b[j];
			}
			b[i] = x[i] - temp;
		}
		
		for(int i = n - 1; i >= 0; i--) {
			temp = 0.0;
			for(int j = i + 1; j < n; j++) {
				temp += x[j] * A[i][j];
			}
			x[i] = (b[i] - temp) / A[i][i];
		}
	}

	/**
	 * 在对矩阵进行列选主元LU分解之上求解线性方程组, 使用该函数前需要使用PLU()函数对矩阵作分解
	 * @param b 线性方程组右端项
	 * @return 方程组解
	 */
	public double[] PLUSolve(double[] b) {
		double[] x = new double[this.n];
		PLUSolve(b, x);
		return x;
	}

	/**
	 * 对称正定矩阵的Cholesky分解
	 */
	public void CholeskyFactorize() {
		for (int k = 0; k < n; k++) {
			A[k][k] = Math.sqrt(A[k][k]);
			for(int j = k + 1; j < n; j++) {
				A[k][j] /= A[k][k];
			}
			for(int i = k + 1; i < n; i++) {
				for(int j = k + 1; j < n; j++) {
					A[i][j] -= A[k][i] * A[k][j];
				}
			}
		}
	}
	
	/**
	 * 
	 * 在对对称正定矩阵进行Cholesky分解之上求解线性方程组
	 * @param b 线性方程组右端项
	 * @param x 方程组解
	 */
	public void CholeskySolve(double[] b, double[] x) {
		Arrays.setAll(x, i -> b[i]);
		double temp = 0.0;
		for(int i = 0; i < n; i++) {
			temp = 0.0;
			for(int j = 0; j < i; j++) {
				temp += A[j][i] * b[j];
			}
			b[i] = (x[i] - temp) / A[i][i];
		}
		
		for(int i = n - 1; i >= 0; i--) {
			temp = 0.0;
			for(int j = i + 1; j < n; j++) {
				temp += x[j] * A[i][j];
			}
			x[i] = (b[i] - temp) / A[i][i];
		}
	}

	/**
	 *
	 * 在对对称正定矩阵进行Cholesky分解之上求解线性方程组
	 * @param b 线性方程组右端项
	 * @return 方程组解
	 */
	public double[] CholeskySolve(double[] b) {
		double[] x = new double[this.n];
		CholeskySolve(b, x);
		return x;
	}
	
	/**
	 * 对称正定矩阵的cholesky分解, 将矩阵分解为A = LDL^T方式.
	 */
	public void CholeskyFactorize2() {
		double[] v = new double[n - 1];
		double temp = 0.0;
		for(int j = 0; j < n; j++) {
			for(int i = 0; i < j; i++) {
				v[i] = A[j][i] * A[i][i];
			}
			for(int i = 0; i < j; i++) {
				A[j][j] -= A[j][i] * v[i];
			}
			

			for(int k = j + 1; k < n; k++) {
				temp = 0.0;
				for(int i = 0; i < j; i++) {
					temp += A[k][i] * v[i];
				}
				A[k][j]  = (A[k][j] - temp) / A[j][j];
			}
		}
	}
	
	/**
	 * 在将正定矩阵分解为LDL^T基础之上求解线性方程组
	 * @param b 方程组右端项
	 * @param x 返回解向量
	 */
	public void CholeskySolve2(double[] b, double[] x) {
		Arrays.setAll(x, i -> b[i]);
		double temp = 0.0;
		//Ly = b
		for(int i = 0; i < n; i++) {
			temp = 0.0;
			for(int j = 0; j < i; j++) {
				temp += A[i][j] * b[j];
			}
			b[i] = x[i] - temp;
		}
		//Dz = y
		Arrays.setAll(b, i -> b[i] / A[i][i]);
		
		//L^Tx = z
		for(int i = n - 1; i >= 0; i--) {
			temp = 0.0;
			for(int j = i + 1; j < n; j++) {
				temp += x[j] * A[j][i];
			}
			x[i] = b[i] - temp;
		}
	}

	/**
	 * 在将正定矩阵分解为LDL^T基础之上求解线性方程组
	 * @param b 方程组右端项
	 * @return  返回解向量
	 */
	public double[] CholeskySolve2(double[] b) {
		double[] x = new double[this.n];
		CholeskySolve2(b, x);
		return x;
	}
	
	int n;
	double[][] A;
	int[] P;
}
